import json

from utils.OpenAiBody import Message, GptQuery
from utils.StaticObj import CHAT_CACHE, GPT_ROLE_SYS, GPT_ROLE_USER, GPT_ROLE_FUN, GPT_ROLE_ASSISTANT


def getParam(username, msgStr, functions, sysStr=None):
    message = Message(GPT_ROLE_USER, msgStr)
    systemMsg = None
    if sysStr:
        systemMsg = Message(GPT_ROLE_SYS, sysStr)
    res = createChat(username, message, systemMessage=systemMsg, functions=functions)
    return res

def getScene(username, msgStr, sysStr=None, functions=None):
    message = Message(GPT_ROLE_USER, msgStr)
    systemMsg = None
    if sysStr:
        systemMsg = Message(GPT_ROLE_SYS, sysStr)
    res = createChat(username, message, systemMessage=systemMsg, functions=functions)
    return res

def getFun(username, data, functionName, sysStr=None, hisMessage=[]):
    datastr = json.dumps(data).replace(' ', '')
    message = Message(GPT_ROLE_FUN, datastr, name=functionName)
    systemMsg = None
    if sysStr:
        systemMsg = Message(GPT_ROLE_SYS, sysStr)
    res = createChat(username, message, systemMessage=systemMsg, hisMessages=hisMessage)
    return res

def getContrast(username, data, functionName, sysStr=None, hisMessage=[]):
    datastr = json.dumps(data).replace(' ', '')
    print('confun',datastr)
    message = Message(GPT_ROLE_FUN, data, name=functionName)
    systemMsg = None
    if sysStr:
        systemMsg = Message(GPT_ROLE_SYS, sysStr)
    res,ec = createContrastChat(username, data, systemMessage=systemMsg, hisMessages=hisMessage)
    return res,ec

def createChat(username, message, hisMessages=[], systemMessage=None, functions=None):
    if username not in CHAT_CACHE:
        CHAT_CACHE[username] = []
    if hisMessages:
        CHAT_CACHE[username] = hisMessages
    CHAT_CACHE[username].append(message)
    print(CHAT_CACHE[username].__len__() > 15)
    if CHAT_CACHE[username].__len__() > 15:
        CHAT_CACHE[username] = CHAT_CACHE[username][15:]
    print("user:" + username, CHAT_CACHE[username].__len__())
    gpt = GptQuery(systemMessage, messages=CHAT_CACHE[username])
    resp = gpt.query(message, functions)
    role = 'assistant'
    content = ''
    function_call = {
        "name": "",
        "arguments": ""
    }
    if gpt.stream:
        choices = []
        for r in resp:
            choices.extend(r.choices)
        for choice in choices:
            if choice['finish_reason'] in ['stop', 'function_call']:
                break
            if 'role' in choice['delta']:
                role = choice['delta']['role']
            if 'content' in choice['delta'] and choice['delta']['content'] is not None:
                content = content + choice['delta']['content']
            if 'function_call' in choice['delta']:
                if 'name' in choice['delta']['function_call']:
                    function_call['name'] = function_call['name'] + choice['delta']['function_call']['name']
                if 'arguments' in choice['delta']['function_call']:
                    function_call['arguments'] = function_call['arguments'] + choice['delta']['function_call']['arguments']
        message = Message(role, content)
        if function_call['name'] != '':
            message.function_call = function_call
        CHAT_CACHE[username].append(message)
        return message
    else:
        print('消耗token:', resp['usage'])
        message = Message(resp['choices'][0]["message"]['role'], resp['choices'][0]["message"]['content'])
        if 'function_call' in resp['choices'][0]["message"]:
            message.function_call = resp['choices'][0]["message"]['function_call']
        CHAT_CACHE[username].append(message)
        return message

def createConChat(username, message, hisMessages=[], systemMessage=None, functions=None):
    print("user:" + username)
    gpt = GptQuery(systemMessage)
    resp = gpt.conquery(message, functions)
    role = 'assistant'
    content = ''
    function_call = {
        "name": "",
        "arguments": ""
    }
    if gpt.stream:
        choices = []
        for r in resp:
            choices.extend(r.choices)
        for choice in choices:
            if choice['finish_reason'] in ['stop', 'function_call']:
                break
            if 'role' in choice['delta']:
                role = choice['delta']['role']
            if 'content' in choice['delta'] and choice['delta']['content'] is not None:
                content = content + choice['delta']['content']
            if 'function_call' in choice['delta']:
                if 'name' in choice['delta']['function_call']:
                    function_call['name'] = function_call['name'] + choice['delta']['function_call']['name']
                if 'arguments' in choice['delta']['function_call']:
                    function_call['arguments'] = function_call['arguments'] + choice['delta']['function_call']['arguments']
        message = Message(role, content)
        if function_call['name'] != '':
            message.function_call = function_call
        return message
    else:
        print('消耗token:', resp['usage'])
        message = Message(resp['choices'][0]["message"]['role'], resp['choices'][0]["message"]['content'])
        if 'function_call' in resp['choices'][0]["message"]:
            message.function_call = resp['choices'][0]["message"]['function_call']
        return message

def createContrastChat(username, message, hisMessages=[], systemMessage=None, functions=None):
    print("user:" + username)
    gpt = GptQuery(systemMessage)
    resp = gpt.contrast(message, functions)
    role = 'assistant'
    content = ''
    ec = ''
    function_call = {
        "name": "",
        "arguments": ""
    }
    if gpt.constream:
        choices = []
        for r in resp:
            choices.extend(r.choices)
        for choice in choices:
            print(choice)
            if choice['finish_reason'] in ['stop', 'function_call']:
                break
            if 'role' in choice['delta']:
                role = choice['delta']['role']
            if 'content' in choice['delta'] and choice['delta']['content'] is not None:
                content = content + choice['delta']['content']
            if 'ec' in choice['delta'] and choice['delta']['ec'] is not None:
                ec = ec + choice['delta']['ec']
            if 'function_call' in choice['delta']:
                if 'name' in choice['delta']['function_call']:
                    function_call['name'] = function_call['name'] + choice['delta']['function_call']['name']
                if 'arguments' in choice['delta']['function_call']:
                    function_call['arguments'] = function_call['arguments'] + choice['delta']['function_call']['arguments']
        message = Message(role, content)
        if function_call['name'] != '':
            message.function_call = function_call
        return message,ec
    else:
        print('ec:', resp['choices'][0]["message"]['content'])
        print('jsonec:', json.loads(resp['choices'][0]["message"]['content']))
        ec = json.loads(resp['choices'][0]["message"]['content'])['ec']
        message = Message(resp['choices'][0]["message"]['role'], json.loads(resp['choices'][0]["message"]['content'])['md'])
        if 'function_call' in resp['choices'][0]["message"]:
            message.function_call = resp['choices'][0]["message"]['function_call']
        return message,ec
